package com.pengxh.app.gb28181.util;

import android.media.Image;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.os.Build;
import android.os.Bundle;
import android.util.Log;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

public class YuvEncoder {
    private static final String TAG = "YuvEncoder";
    private static final String MIME_TYPE = "video/avc"; // H.264编码器——硬编码
    private MediaCodec mMediaCodec;
    private final int mWidth;
    private final int mHeight;
    private boolean mIsEncoding = false;
    private static final int QUEUE_CAPACITY = 30; // 队列容量，根据需求调整
    private final BlockingQueue<ImageFrame> mFrameQueue = new ArrayBlockingQueue<>(QUEUE_CAPACITY);
    private EncoderThread mEncoderThread;

    public interface H264EncodeCallback {
        void onSpsPpsDecoded(byte[] sps, byte[] pps);

        /**
         * @param h264           编码后的H.264数据。这就是原始 H.264 NALU
         * @param timestamp90kHz 90kHz时间戳
         * @param isKeyFrame     是否I帧
         */
        void onYuvDataEncoded(byte[] h264, long timestamp90kHz, boolean isKeyFrame);
    }

    private final H264EncodeCallback mEncodeCallback;

    // 帧数据封装类
    private static class ImageFrame {
        byte[] data;
        long timestamp90kHz;

        ImageFrame(byte[] data, long timestamp90kHz) {
            this.data = data;
            this.timestamp90kHz = timestamp90kHz;
        }
    }

    public YuvEncoder(int width, int height, H264EncodeCallback callback) {
        mWidth = width;
        mHeight = height;
        mEncodeCallback = callback;
    }

    public void prepare() {
        Log.d(TAG, "prepare: 初始化编码器");
        try {
            mMediaCodec = MediaCodec.createEncoderByType(MIME_TYPE);

            int bitRate;
            int frameRate = 25; // GB28181 常用25fps
            if (mWidth >= 1920 || mHeight >= 1080) {
                // 1080p 或更高
                bitRate = 3000000; // 3.0 Mbps
            } else if (mWidth >= 1280 || mHeight >= 720) {
                // 720p
                bitRate = 2000000; // 2.0 Mbps
            } else {
                // VGA 或更低（如 640x480）
                bitRate = 1200000; // 1.2 Mbps
                frameRate = 20;    // 可适当降低帧率
            }

            MediaFormat mMediaFormat = MediaFormat.createVideoFormat(MIME_TYPE, mWidth, mHeight);
            mMediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar); // NV12对应的格式
            mMediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, bitRate);
            mMediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, frameRate);
            mMediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1); // 每1秒一个I帧
            mMediaFormat.setInteger(MediaFormat.KEY_PROFILE, MediaCodecInfo.CodecProfileLevel.AVCProfileBaseline);
            mMediaFormat.setInteger(MediaFormat.KEY_LEVEL, MediaCodecInfo.CodecProfileLevel.AVCLevel31);
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
                mMediaFormat.setInteger(MediaFormat.KEY_MAX_B_FRAMES, 0);      // 关闭B帧
            }

            mMediaCodec.configure(mMediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
            mMediaCodec.start();

            // 请求立即生成一个 IDR 帧（关键帧）
            Bundle params = new Bundle();
            params.putInt(MediaCodec.PARAMETER_KEY_REQUEST_SYNC_FRAME, 0);
            mMediaCodec.setParameters(params);

            mIsEncoding = true;
            mEncoderThread = new EncoderThread();
            mEncoderThread.start();

            Log.d(TAG, "编码器初始化成功");
        } catch (Exception e) {
            Log.e(TAG, "prepare: ", e);
            release();
        }
    }

    public void encodeFrame(Image image, long timestamp90kHz) {
        if (!mIsEncoding || mMediaCodec == null) {
            Log.w(TAG, "Encoder not prepared");
            return;
        }

        try {
            byte[] nv12 = ImageConverter.convertToNV12(image);
            boolean offered = mFrameQueue.offer(new ImageFrame(nv12, timestamp90kHz));
            if (!offered) {
                String error = "编码队列已满，丢弃一帧数据";
                Log.w(TAG, error);
            }
        } catch (Exception e) {
            Log.e(TAG, "encodeFrame: ", e);
        } finally {
            if (image != null) {
                image.close();
            }
        }
    }

    /**
     * 编码线程，负责从队列中取出数据并进行编码
     */
    @SuppressWarnings("all")
    private class EncoderThread extends Thread {
        @Override
        public void run() {
            while (mIsEncoding && !isInterrupted()) {
                try {
                    // 从队列中获取待编码数据
                    ImageFrame imageFrame = mFrameQueue.take();
                    if (imageFrame == null) {
                        continue;
                    }

                    // 获取输入缓冲区，计算微秒级时间戳，用于MediaCodec
                    long presentationTimeUs = imageFrame.timestamp90kHz * 1000000L / 90000L; // 转为 us
                    int inputBufferIndex = mMediaCodec.dequeueInputBuffer(-1);
                    if (inputBufferIndex >= 0) {
                        ByteBuffer inputBuffer = mMediaCodec.getInputBuffer(inputBufferIndex);
                        if (inputBuffer != null) {
                            inputBuffer.clear();
                            inputBuffer.put(imageFrame.data);
                            // 将数据送入编码器
                            mMediaCodec.queueInputBuffer(inputBufferIndex, 0, inputBuffer.limit(), presentationTimeUs, 0);
                        }
                    }

                    // 获取编码后的输出数据
                    MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
                    int outputBufferIndex = mMediaCodec.dequeueOutputBuffer(bufferInfo, 0);
                    while (outputBufferIndex >= 0) {
                        // 处理编码后的H.264数据
                        ByteBuffer outputBuffer = mMediaCodec.getOutputBuffer(outputBufferIndex);
                        if (outputBuffer != null && bufferInfo.size > 0) {
                            // 检查是否包含编解码器配置信息
                            if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0) {
                                // 保存 SPS/PPS（用于 SDP 上报，不用于媒体流）
                                byte[] configBytes = new byte[bufferInfo.size];
                                outputBuffer.get(configBytes);

                                // 解析SPS和PPS数据
                                List<byte[]> nalus = splitNalu(configBytes);
                                byte[] sps = null, pps = null;
                                for (byte[] nalu : nalus) {
                                    int type = nalu[0] & 0x1F;
                                    if (type == 7) sps = nalu;
                                    else if (type == 8) pps = nalu;
                                }
                                if (sps != null && pps != null) {
                                    mEncodeCallback.onSpsPpsDecoded(sps, pps);
                                }
                            } else {
                                // 普通视频帧：只含 slice NALU（带起始码）
                                byte[] h264Bytes = new byte[bufferInfo.size];
                                outputBuffer.get(h264Bytes);

                                boolean isKeyFrame = (bufferInfo.flags & MediaCodec.BUFFER_FLAG_KEY_FRAME) != 0;

                                List<byte[]> nalus = splitNalu(h264Bytes); // 分割出裸 NALU
                                for (byte[] nalu : nalus) {
                                    if (nalu.length == 0) continue;
                                    int naluType = nalu[0] & 0x1F;
                                    // 只推送 slice
                                    if (naluType == 1 || naluType == 5 || naluType == 19) {
                                        mEncodeCallback.onYuvDataEncoded(nalu, imageFrame.timestamp90kHz, isKeyFrame);
                                    }
                                }
                            }
                        }
                        mMediaCodec.releaseOutputBuffer(outputBufferIndex, false);
                        outputBufferIndex = mMediaCodec.dequeueOutputBuffer(bufferInfo, 0);
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                    Log.d(TAG, "编码线程被中断");
                    mIsEncoding = false;
                }
            }
        }
    }

    private static List<byte[]> splitNalu(byte[] frame) {
        List<byte[]> naluList = new ArrayList<>();
        int index = 0;
        while (index < frame.length) {
            // 查找 0x000001 或 0x00000001
            if (index + 4 <= frame.length &&
                    frame[index] == 0 && frame[index + 1] == 0 &&
                    frame[index + 2] == 0 && frame[index + 3] == 1) {
                int start = index + 4;
                int end = findNextStartCode(frame, start);
                naluList.add(Arrays.copyOfRange(frame, start, end));
                index = end;
            } else if (index + 3 <= frame.length &&
                    frame[index] == 0 && frame[index + 1] == 0 &&
                    frame[index + 2] == 1) {
                int start = index + 3;
                int end = findNextStartCode(frame, start);
                naluList.add(Arrays.copyOfRange(frame, start, end));
                index = end;
            } else {
                index++;
            }
        }
        return naluList;
    }

    private static int findNextStartCode(byte[] data, int start) {
        for (int i = start; i < data.length - 3; i++) {
            if ((data[i] == 0 && data[i + 1] == 0 && data[i + 2] == 0 && data[i + 3] == 1) ||
                    (data[i] == 0 && data[i + 1] == 0 && data[i + 2] == 1)) {
                return i;
            }
        }
        return data.length;
    }

    public void release() {
        mIsEncoding = false;

        if (mEncoderThread != null) {
            mEncoderThread.interrupt();
            try {
                mEncoderThread.join();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            mEncoderThread = null;
        }

        if (mMediaCodec != null) {
            mMediaCodec.stop();
            mMediaCodec.release();
            mMediaCodec = null;
        }

        mFrameQueue.clear();
        Log.d(TAG, "编码器已释放");
    }
}
